namespace IBS.Tests

open IBS.Math
open FsUnit
open NUnit.Framework

module TvmTest =
    [<Test>]
    let ``FV of investment with monthly compounding``() =
        let fvalue = fv 0.05 10. -500. -5000. Frequency.Monthly
        fvalue |> should approxEqual 85876.18721129
        
    [<Test>]
    let ``PV of investment with monthly compounding``() =
        let pvalue = pv 0.05 10. -500. 50000. Frequency.Monthly
        pvalue |> should approxEqual 16782.62314917
        
    [<Test>]
    let ``PV of FV should be equal to initial investment``() =
        let fvalue = fv 0.05 10. -500. -5000. Frequency.Monthly
        let pvalue = pv 0.05 10. -500. fvalue Frequency.Monthly
        pvalue |> should approxEqual -5000.

    [<Test>]
    let ``PV of cash stream to PV of fv + PV of payments``() =
        let pvalue = pv 0.05 10. -500. -5000. Frequency.Semiannually
        let pvPmts = pv 0.05 10. -500. 0. Frequency.Semiannually
        let pvFV   = pv 0.05 10. 0. -5000. Frequency.Semiannually
        pvalue |> should approxEqual (pvFV + pvPmts)

    [<Test>]
    let ``find a number of periods needed for $10k -> $50k``() =
        let years = nper 0.05 -500. -10000. 50000.
        years |> should approxEqual 22.51708531

    [<Test>]
    let ``compute monthly payments``() =
        let payment = pmt 0.05 10. -5000. 50000. Frequency.Monthly
        payment |> should approxEqual -268.96148524

    [<Test>]
    let ``find an approximate rate corresponding the investment parameters``() =
        let r = rate 10. -500. -10000. 50000.
        r |> should approxEqual 0.148477940

                
module ForecastingTest =
    open Forecasting
    
    /// Attention: probabilities should match the vars set
    let factors() =
        //f1 contains P(X_1)
        let f1 = { Vars = set [ 1, 2 ]; Val = [| 0.11; 0.89 |] }
        //f2 contains P(X_2 | X_1)
        let f2 = { Vars = set [ 1, 2; 2, 2 ]; Val = [| 0.59; 0.22; 0.41; 0.78 |] }
        //f3 contains P(X_3 | X_2)
        let f3 = { Vars = set [ 2, 2; 3, 2]; Val = [| 0.39; 0.06; 0.61; 0.94 |] }
        f1, f2, f3

    [<Test>]
    let ``factors product``() =
        let f1, f2, _ = factors()
        let fproduct = f1 * f2
        fproduct.Vars |> should equal (set [1, 2; 2, 2])
        fproduct.Val |> should equal [| 0.0649; 0.1958; 0.0451; 0.6942 |]
    
    [<Test>]
    let ``factors marginalization``() =
        let _, f2, _ = factors()
        let fmargin = f2.Marginalize [ 2 ]
        fmargin.Vars |> should equal (set [1,2])
        fmargin.Val |> should equal [|1.; 1.|]

    [<Test>]
    let ``evidence f1``() = 
        let f1, _, _ = factors()
        let fevidence = f1.ApplyEvidence (set [2, 0; 3, 1])
        fevidence.Vars |> should equal (set [1, 2])
        fevidence.Val |> should equal [|0.11; 0.89|]
        
    [<Test>]
    let ``evidence f2``() =
        let _, f2, _ = factors()
        let fevidence = f2.ApplyEvidence (set [2, 0; 3, 1]) 
        fevidence.Vars |> should equal (set [1, 2; 2, 2])
        fevidence.Val |> should equal [|0.59; 0.22; 0.0; 0.0|]

    [<Test>]
    let ``evidence f3``() =
        let _, _, f3 = factors()
        let fevidence = f3.ApplyEvidence (set [2, 0; 3, 1]) 
        fevidence.Vars |> should equal (set [2, 2; 3, 2])
        fevidence.Val |> should equal [|0.0; 0.0; 0.61; 0.0|]
        
    [<Test>]
    let ``joint distribution``() =
        let f1, f2, f3 = factors()
        let fjoint = Seq.reduce ( * ) [f1; f2; f3]
        fjoint.Vars |> should equal (set [1, 2; 2, 2; 3, 2])
        fjoint.Val |> should equal [|0.025311; 0.076362; 0.002706; 0.041652; 0.039589; 0.119438; 0.042394; 0.652548|]
    
    [<Test>]
    let ``marginal``() =
        let f1, f2, f3 = factors()
        let fmarginal = Factor.ComputeMarginal [2; 3] [f1;f2;f3] (set [1, 1])
        fmarginal.Vars |> should equal (set [2, 2; 3, 2])
        fmarginal.Val |> should equal [|0.0858; 0.0468; 0.1342; 0.7332|]
     
    [<Test>]
    let ``marginalization over forecast model``() =
        let x = Model().MarginalizeOver [8] [1, int Age.Below22; 5, int Income.Medium; 7, int FutureIncome.Promising]
        x.Vars |> should equal (set [8, 2])
        x.Val |> should equal [| 0.5926667; 0.4073333 |]
    
    [<Test>]
    let ``forecast``() =
        let f = Model().Forecast(age = Age.Below22, income = Income.Medium, futureIncome = FutureIncome.Promising)
        let probas = Array.sortBy snd f |> Array.map fst
        probas |> should equal [| 0.5926667; 0.4073333 |]
